import os
import tarfile

from exceptions import ModelExtractionException

def write_named_pdbs( path, names, dirs, tarballs ):
	try:
		f = open( path, 'w' )
	except Exception as e:
		raise ModelExtractionException( "Failure opening output file: %s" % (e) )

	# cache the full paths of the tarball contents, as well as the filenames
	tarball_pdb_names,tarball_pdb_paths = [],[]
	for t in tarballs:
		tarball_pdb_paths.append( t.getnames() )
		tarball_pdb_names.append( map(os.path.basename, tarball_pdb_paths[-1]) )

	for (i,n) in enumerate(names):
		tmp, pdb = "%s.pdb" % n, None

		for d in dirs:
			path = os.path.join(d,tmp)

			if( os.access( path, os.R_OK ) ):
				if( pdb != None ):
					print "WARNING: duplicate named PDBs (\"%s\") detected!" % (n)
				else:
					try:
						pdb = open(path, 'r').read().strip()
					except Exception as e:
						raise ModelExtractionException( "Could not read PDB at \"%s\"" % (path) )

		for (j,t) in enumerate(tarballs):
			try: # attempt to find the position of the pdb in the tarball
				k = tarball_pdb_names[j].index(tmp)
			except:
				continue

			if( pdb != None ):
				print "WARNING: duplicate named PDBs (\"%s\") detected!" % (n)
			else:
				try:
					member = t.getmember( tarball_pdb_paths[j][k] )
					pdb = t.extractfile(member).read().strip()
				except Exception as e:
					raise ModelExtractionException( "Could not extract PDB \"%s\" from tarball \"%s\"" % (n,t.name) )

		if(pdb == None):
			raise ModelExtractionException( "Could not find pdb \"%s\" in any of the provided resources" % (n) )

		try:
			f.write("MODEL        %d\n%s\nENDMDL\n" % (i+1, pdb ) )
		except Exception as e:
			raise ModelExtractionException( "Could not append PDB to output: %s" % (e) )

	f.close()

def write_model_attributes( path, name, attributes ):
	f = open( path, 'w' )
	f.write("""# generated by make_generation_models
attribute: %s
match mode: 1-to-1
recipient: molecules
""" % name)
	for (i,a) in enumerate(attributes):
		f.write("\t#0.%i\t%.3f\n" % (i+1,a))
	f.close()

def write_model_script( path, name, attributes ):
	# this is some meta shit. Could probably achieve the same thing with pickle or serialize, but eh

	values = "[%s]" % (','.join([str(a/max(attributes)) for a in attributes]))

	script = """
from colorsys	import rgb_to_hsv,hsv_to_rgb
from pymol		import cmd

class MESMER_colors:

	def __init__(self, name= '', values=[]):
		self.values = values
		self.name = name
		self.make_colors()

		cmd.extend("<<NAME>>_color", lambda x='blue',y='red' : self.make_colors(x,y) )
		cmd.extend("<<NAME>>_apply", lambda x='*',y='cartoon_color' : self.apply_colors(x,y) )

	def _interpolate_color(self, value, start, end ):
		rA,rB = cmd.get_color_tuple(start), cmd.get_color_tuple(end)
		hA,hB = rgb_to_hsv( rA[0], rA[1], rA[2] ), rgb_to_hsv( rB[0], rB[1], rB[2] )

		def transition( v, s, e ):
			return s + (e-s)*v

		def transition3( v, (s1,s2,s3), (e1,e2,e3) ):
			return (transition(v,s1,e1), transition(v,s2,e2), transition(v,s3,e3))

		ret = transition3(value,hA,hB)
		return hsv_to_rgb( ret[0], ret[1], ret[2] )

	def make_colors(self, start_color='blue', end_color='red' ):
		self.colors = [None]*len(self.values)
		for (i,w) in enumerate(self.values):
			self.colors[i] = "Mc_%s_%i" % (self.name,i)
			cmd.set_color( "Mc_%s_%i" % (self.name,i), self._interpolate_color(w, start_color, end_color) )

	def apply_colors(self, model='<<NAME>>', repr="cartoon_color" ):
		cmd.set("all_states", "on")
		for (i,c) in enumerate(self.colors):
			cmd.set(repr,c,model,i)

<<NAME>>_colors = MESMER_colors( '<<NAME>>', <<VALUES>> )
print "Type <<NAME>>_color to set color gradient, e.g. <<NAME>>_color blue, red"
print "Type <<NAME>>_apply to apply color to model depiction, e.g. <<NAME>>_apply <<NAME>>, cartoon_color"
"""
	script = script.replace('<<VALUES>>',values)
	script = script.replace('<<NAME>>',name)

	f = open( path, 'w' )
	f.write( script )
	f.close()

